# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging
from collections import namedtuple

import fibers
from absl import app, flags

from coin.base.datetime_util import to_datetime
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.subscription import SubscriptionRequest, FeedSubscriptionRequest
from coin.feed.fastfeed.feed_cache import run_from_feed_cache
from coin.feed.fastfeed.util import (
    get_products_from_sub_req,
)

Message = namedtuple('Message', ['is_book', 'product', 'timestamp', 'book', 'trade'])


def dump_levels(book_iter, product_holder=None, max_levels=None):
  max_levels = max_levels or 100000
  levels = []
  for level in book_iter:
    if len(levels) == max_levels:
      break

    if product_holder:
      price = product_holder.sanitize_price(level.price)
      qty = product_holder.sanitize_qty(level.qty)
    else:
      price = level.price
      qty = level.qty

    if qty == 0:
      continue

    levels.append((price, qty))
  return levels


class FeedCacheDiffer:
  def __init__(self,
               machine,
               date,
               worker_id,
               feed_cache_dir,
               feed_sub_req=None,
               sub_req=None,
               products=None,
               check_products=None):
    self.machine = machine
    self.worker_id = worker_id
    self.feed_cache_dir = feed_cache_dir

    if feed_sub_req:
      self.feed_sub_req = feed_sub_req
    else:
      self.products = products or get_products_from_sub_req(sub_req, date)
      self.feed_sub_req = FeedSubscriptionRequest.create(self.products, sub_req)

    all_products = []
    for _, prods in self.feed_sub_req.gen_iter():
      all_products += prods
    self.check_products = set(check_products or all_products)

    self.from_ts = to_datetime(date)
    self.to_ts = self.from_ts + datetime.timedelta(days=1)

    self._archive_msg = Message(False, None, 0, None, None)
    self._feed_cache_msg = Message(False, None, 0, None, None)

    self._fiber_archive = fibers.Fiber(self.run_archive)
    self._fiber_feed_cache = fibers.Fiber(self.run_feed_cache)

  def run(self):
    self._fiber_archive.switch()

  def run_archive(self):
    run_from_archive(self.feed_sub_req,
                     self.on_archive_book_reset,
                     self.from_ts,
                     self.to_ts,
                     machine=self.machine,
                     worker_id=self.worker_id)

  def run_feed_cache(self):
    run_from_feed_cache(self.feed_sub_req,
                        self.on_feed_cache_book_reset,
                        self.from_ts,
                        self.to_ts,
                        machine=self.machine,
                        worker_id=self.worker_id,
                        feed_cache_dir=self.feed_cache_dir)

  def compare(self):
    if self._archive_msg.timestamp != self._feed_cache_msg.timestamp:
      print(to_datetime(self._archive_msg.timestamp),
            to_datetime(self._feed_cache_msg.timestamp),
            'timestamp',
            self._archive_msg,
            self._feed_cache_msg)
      return

    if self._archive_msg.product != self._feed_cache_msg.product:
      print(to_datetime(self._archive_msg.timestamp),
            'product %s != %s' % (self._archive_msg.product, self._feed_cache_msg.product))
      return

    if self._archive_msg.is_book != self._feed_cache_msg.is_book:
      print(to_datetime(self._archive_msg.timestamp),
            'is_book',
            self._archive_msg,
            self._feed_cache_msg)
      return

    product_holder = get_holder_from_product(self._archive_msg.product)
    assert product_holder is not None

    if self._archive_msg.is_book:
      archive_ask = dump_levels(self._archive_msg.book.iter_ask(), product_holder, 20)
      feed_cache_ask = dump_levels(self._feed_cache_msg.book.iter_ask(), None, 20)
      if archive_ask != feed_cache_ask:
        print(to_datetime(self._archive_msg.timestamp),
              self._archive_msg.product,
              'ask book\n%s\n%s' % (archive_ask, feed_cache_ask))

      archive_bid = dump_levels(self._archive_msg.book.iter_bid(), product_holder, 20)
      feed_cache_bid = dump_levels(self._feed_cache_msg.book.iter_bid(), None, 20)
      if archive_bid != feed_cache_bid:
        print(to_datetime(self._archive_msg.timestamp),
              self._archive_msg.product,
              'bid book\n%s\n%s' % (archive_bid, feed_cache_bid))

    if not self._archive_msg.is_book:
      atrade = self._archive_msg.trade
      btrade = self._feed_cache_msg.trade
      if (product_holder.sanitize_price(atrade.price) != btrade.price
          or product_holder.sanitize_qty(atrade.qty) != btrade.qty or atrade.side != btrade.side):
        print(to_datetime(self._archive_msg.timestamp), 'trade %s != %s' % (atrade, btrade))

  def on_archive_book_reset(self, sub_req, feed_handler):
    logging.info('Archive book reset')
    for product in self.feed_sub_req.get_products(sub_req):
      feed_handler.subscribe_book(product, functools.partial(self.on_archive_book, product))
      feed_handler.subscribe_trade(product, functools.partial(self.on_archive_trade, product))

  def on_archive_book(self, product, book):
    if product not in self.check_products:
      return
    if self._archive_msg.timestamp < self._feed_cache_msg.timestamp:
      self._archive_msg = Message(True, product, book.timestamp, book, None)
      self.compare()
      return
    self._archive_msg = Message(True, product, book.timestamp, book, None)
    self._fiber_feed_cache.switch()

  def on_archive_trade(self, product, trade):
    if product not in self.check_products:
      return
    product_holder = get_holder_from_product(product)
    assert product_holder is not None

    if product_holder.sanitize_qty(trade.qty) == 0:
      logging.debug(trade)
      return

    if self._archive_msg.timestamp < self._feed_cache_msg.timestamp:
      self._archive_msg = Message(False, product, trade.timestamp, None, trade)
      self.compare()
      return

    self._archive_msg = Message(False, product, trade.timestamp, None, trade)
    self._fiber_feed_cache.switch()

  def on_feed_cache_book_reset(self, sub_req, feed_handler):
    logging.info('Feed cache book reset')
    for product in self.feed_sub_req.get_products(sub_req):
      feed_handler.subscribe_book(product, functools.partial(self.on_feed_cache_book, product))
      feed_handler.subscribe_trade(product, functools.partial(self.on_feed_cache_trade, product))

  def on_feed_cache_book(self, product, book):
    if product not in self.check_products:
      return
    self._feed_cache_msg = Message(True, product, book.timestamp, book, None)
    self.compare()
    if self._archive_msg.timestamp > self._feed_cache_msg.timestamp:
      return
    self._fiber_archive.switch()

  def on_feed_cache_trade(self, product, trade):
    if product not in self.check_products:
      return
    self._feed_cache_msg = Message(False, product, trade.timestamp, None, trade)
    self.compare()
    if self._archive_msg.timestamp > self._feed_cache_msg.timestamp:
      return
    self._fiber_archive.switch()


def main(argv):
  FLAGS = flags.FLAGS
  date = datetime.datetime.strptime(FLAGS.date, '%Y%m%d').date()
  '''
  from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
  from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
  feed_sub_req = FeedSubscriptionRequest()
  feed_sub_req.add_products([BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')],
                            SubscriptionRequest('Futures', 'Bitmex'))
  feed_sub_req.add_products([BitflyerFuturesProduct.FromStr('BTC-JPY.PERPETUAL')],
                            SubscriptionRequest('Futures', 'Bitflyer'))

  differ = FeedCacheDiffer(FLAGS.machine, date, FLAGS.worker_id, FLAGS.feed_cache_dir,
                           feed_sub_req=feed_sub_req)
  differ.run()
  '''

  sub_req = SubscriptionRequest.from_str(FLAGS.sub_req)

  # Filters products
  all_products = get_products_from_sub_req(sub_req, date, FLAGS.group)
  products = []
  for p in all_products:
    if ((not FLAGS.base_filter or p.base.currency in FLAGS.base_filter)
        and (not FLAGS.quote_filter or p.quote.currency in FLAGS.quote_filter)):
      products.append(p)
    elif (FLAGS.symbol_filter and p.symbol in FLAGS.symbol_filter):
      products.append(p)

  for product in products:
    logging.info('%s: %s', sub_req, product)
    differ = FeedCacheDiffer(FLAGS.machine,
                             date,
                             FLAGS.worker_id,
                             FLAGS.feed_cache_dir,
                             sub_req=sub_req,
                             products=[product])
    differ.run()


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_string('machine', 'feed-01.ap-northeast-1.aws', 'machine')
  flags.DEFINE_string('date', None, 'Date')
  flags.DEFINE_string(
      'sub_req',
      'Futures.Bitmex',
      'Subscription Request String. <market_type>.<exchange>.<api_version> or '
      '<market_type>.<exchange>')
  flags.DEFINE_string('worker_id', '1', 'Worker ID')
  flags.DEFINE_string('feed_cache_dir', '~/data/feed_cache', 'Feed cache directory')

  flags.DEFINE_list('base_filter',
                    [
                        'BTC',
                        'ETH',
                        'LTC',
                        'ETC',
                        'BCHABC',
                        'XRP',
                        'EOS',
                        'BNB',
                        'ADA',
                        'BCHSV',
                        'TRX',
                        'LINK',
                        'XLM',
                        'STEEM',
                        'TMTG',
                        'FCT',
                        'BHP',
                        'THETA',
                        'RNT',
                        'FNB',
                        'INS',
                        'META',
                        'EL',
                        'TON',
                        'SXP',
                    ],
                    'Quote currency filter')
  flags.DEFINE_list('quote_filter',
                    ['USD', 'KRW', 'JPY', 'USDT', 'USDC', 'TUSD', 'PAX', 'USDK', 'BTC', 'ETH'],
                    'Base curency filter')
  flags.DEFINE_list('symbol_filter', ['BTC-HUSD'], 'symbol filter')
  flags.DEFINE_string('group', 'all', 'Empty group means all groups.')
  app.run(main)
